# 1. 导入必要库
import torch
import torch.distributed as dist
from diffusers import StableVideoDiffusionPipeline
from PIL import Image
import requests
from io import BytesIO
import psutil
from moviepy import ImageSequenceClip


# 2. 初始化分布式环境
def init_ddp():
    dist.init_process_group(
        backend='gloo',  # Windows推荐使用gloo后端
        init_method='tcp://192.168.1.100:29500',  # 主节点IP和端口
        world_size=4,  # 总进程数（建议等于GPU数）
        rank=0  # 当前进程编号（0为主节点）
    )


# 3. 分布式帧生成函数
@torch.no_grad()
def generate_frames_distributed(init_image, num_frames, device):
    # 加载模型（仅主进程下载）
    if dist.get_rank() == 0:
        pipeline = StableVideoDiffusionPipeline.from_pretrained(
            "stabilityai/stable-video-diffusion-img2vid-xt",
            device=device,
            torch_dtype=torch.float16,
            cache_dir='models'
        ).to(device)
        # 广播模型到其他进程
        dist.broadcast_object_list([pipeline], src=0)
    else:
        pipeline = dist.broadcast_object_list([None], src=0)[0]
    # 启用xFormers加速注意力
    pipeline.enable_xformers_memory_efficient_attention()
    # 梯度检查点（减少峰值显存）
    pipeline.unet.enable_gradient_checkpointing()
    # 分配任务（每进程生成2帧）
    frame_indices = list(range(num_frames))
    chunk_size = num_frames // dist.get_world_size()
    local_indices = frame_indices[dist.get_rank() * chunk_size: (dist.get_rank() + 1) * chunk_size]

    # 生成视频帧
    local_frames = []
    for idx in local_indices:
        with torch.cuda.amp.autocast():
            frame = pipeline(init_image).frames[0]
        local_frames.append(frame)

    # 收集所有帧到主进程
    all_frames = [[] for _ in range(dist.get_world_size())]
    dist.all_gather_object(all_frames[dist.get_rank()], local_frames)

    # 主进程合并结果
    if dist.get_rank() == 0:
        merged_frames = [frame for chunk in all_frames for frame in chunk]
        return merged_frames
    return None


# 4. 主函数
def main():
    # 初始化DDP
    init_ddp()
    device = torch.device(f'cuda:{dist.get_rank()}')

    # 加载初始图像（主进程下载）
    if dist.get_rank() == 0:
        url = "https://pic.rmb.bdstatic.com/bjh/other/7132183c3a3973c03796a833e857df67.jpeg"
        response = requests.get(url)
        init_image = Image.open(BytesIO(response.content)).convert("RGB")
        # 广播图像到其他进程
        dist.broadcast_object_list([init_image], src=0)
    else:
        init_image = dist.broadcast_object_list([None], src=0)[0]

    # 生成视频帧
    video_frames = generate_frames_distributed(init_image, 8, device)

    # 主进程保存结果
    if dist.get_rank() == 0:
        clip = ImageSequenceClip(video_frames, fps=4)
        clip.write_videofile("output.mp4")

    # 清理资源
    dist.destroy_process_group()


if __name__ == "__main__":
    main()